package karpiel.boguchwal.wojciech.przemek

import karpiel.boguchwal.wojciech.przemek.Ast._

object Utils {
  def shift(expr: Ast, minToAffect: Int = 0, delta: Int): Ast = expr match {
    case PiType(typeT, value) => PiType(shift(typeT, minToAffect, delta), shift(value, minToAffect + 1, delta))
    case PiValue(typeT, value) => PiValue(shift(typeT, minToAffect, delta), shift(value, minToAffect + 1, delta))
    case Bruijn(i, name) => Bruijn(if (i >= minToAffect) delta + i else i, name)
    //to się może stać przy bruinizacji, gdzie indziej jest błędem
    case NamedType(name, typeT) => NamedType(name, shift(typeT, minToAffect, delta))
    case Var(name) => Var(name)
    //
    case any => Utils.applyToSubTrees(any, e => shift(e, minToAffect, delta))
  }

  def unshift(expr: Ast, minToAffect: Int = 0, delta: Int): Ast = shift(expr, minToAffect, -delta)

  /**
    * TODO ogarnąć zamianę (pi x (pi ax expr)) ->
    *
    *
    * 1. zmienić (br (1+ n)) na (pr1 (br n))
    * 2. zmienić (br n) na (pr2 (br n))
    * 3. zmianień (nr n) na (br (1- n))
    *
    * @return
    */
  def sigmaElimHelper(expr: Ast, index: Int): Ast = expr match {
    case Bruijn(i, name) if i == index + 1 => Proj1(Bruijn(index))
    case Bruijn(i, name) if i == index => Proj2(Bruijn(index))
    case Bruijn(i, name) if i > index + 1 => Bruijn(i - 1)
    case Bruijn(i, name) if i < index /*else*/ => Bruijn(i, name)
    case PiValue(typeT, value) =>
      PiValue(typeT.sigmaElimHelper(index), value.sigmaElimHelper(index + 1))
    case PiType(typeT, value) =>
      PiType(typeT.sigmaElimHelper(index), value.sigmaElimHelper(index + 1))
    case any => Utils.applyToSubTrees(any, e => sigmaElimHelper(e, index))
  }

  def applyToSubTrees(expr: Ast, fn: Ast => Ast): Ast = {
    expr match {
      case UnitValue => UnitValue
      case UnitType => UnitType
      case Universe(i) => Universe(i)
      case Bruijn(i, n) => Bruijn(i, n)
      case PiType(typeT, value) => PiType(fn(typeT), fn(value))
      case PiValue(typeT, value) => PiValue(fn(typeT), fn(value))
      case PiElim(piValue, argument) => PiElim(fn(piValue), fn(argument))
      case SumType(aType, bType) => SumType(fn(aType), fn(bType))
      case SumInl(aValue, bType) => SumInl(fn(aValue), fn(bType))
      case SumInr(aType, bValue) => SumInr(fn(aType), fn(bValue))
      case SumElim(aFn, bFn) => SumElim(fn(aFn), fn(bFn))
      case SigmaType(pi) => SigmaType(fn(pi))
      case SigmaValue(x, aX, a) => SigmaValue(fn(x), fn(aX), fn(a))
      case SigmaElim(pi) => SigmaElim(fn(pi))
      case Proj1(sigma) => Proj1(fn(sigma))
      case Proj2(sigma) => Proj2(fn(sigma))
      case IdType(x, y) => IdType(fn(x), fn(y))
      case Refl(x) => Refl(fn(x))
      case J(typeFunc, reflFunc) => J(fn(typeFunc), fn(reflFunc))
      case BoolType => BoolType
      case BoolTrue => BoolTrue
      case BoolFalse => BoolFalse
      case BoolElim(ifTrue, ifFalse) => BoolElim(fn(ifTrue), fn(ifFalse))
      case BoolRec(a, b, c) => BoolRec(a | fn, b | fn, c | fn)
      case EmptyType => EmptyType
      case EmptyElim(anyType) => EmptyElim(fn(anyType))
      case WType(arityFn) => WType(arityFn | fn)
      case WElim(e, fn_) => WElim(e | fn, fn_ | fn)
      case WSup(wType, label, fn_) => WSup(wType | fn, label | fn, fn_ | fn)
      case Var(name) => Var(name)
      case NamedType(name, typeT) => NamedType(name, fn(typeT))
    }
  }

  def usesBruins(expr: Ast, delta: Int = 0): Set[Int] = {
    expr match {
      case Bruijn(i, _) => Set(i + delta).filter(_ >= 0)
      case Ast.PiValue(typeT, value) => (usesBruins(typeT, delta) ++ usesBruins(value, delta - 1)).filter(_ >= 0)
      case Ast.PiType(typeT, value) => (usesBruins(typeT, delta) ++ usesBruins(value, delta - 1)).filter(_ >= 0)
      case EmptyType => Set()
      case BoolFalse => Set()
      case BoolTrue => Set()
      case BoolType => Set()
      case Universe(_) => Set()
      case UnitType => Set()
      case UnitValue => Set()
      case q => applyToSubTrees[Set[Int]](q, usesBruins(_, delta), _.flatten.toSet.filter(_ >= 0))
    }

  }

  def applyToSubTrees[T](expr: Ast, fn: Ast => T, aggregator: Seq[T] => T): T = {
    def aggregator_(args: T*) = aggregator(args.toSeq)

    expr match {
      case UnitValue => UnitValue | fn | ((a: T) => aggregator_(a))
      case UnitType => UnitType | fn | (a => aggregator_(a))
      case Universe(i) => Universe(i) | fn | (a => aggregator_(a))
      case Bruijn(i, n) => Bruijn(i, n) | fn | (a => aggregator_(a))
      case PiType(typeT, value) => aggregator_(fn(typeT), fn(value))
      case PiValue(typeT, value) => aggregator_(fn(typeT), fn(value))
      case PiElim(piValue, argument) => aggregator_(fn(piValue), fn(argument))
      case SumType(aType, bType) => aggregator_(fn(aType), fn(bType))
      case SumInl(aValue, bType) => aggregator_(fn(aValue), fn(bType))
      case SumInr(aType, bValue) => aggregator_(fn(aType), fn(bValue))
      case SumElim(aFn, bFn) => aggregator_(fn(aFn), fn(bFn))
      case SigmaType(pi) => aggregator_(fn(pi))
      case SigmaValue(x, aX, a) => aggregator_(fn(x), fn(aX), fn(a))
      case SigmaElim(pi) => aggregator_(fn(pi))
      case Proj1(sigma) => aggregator_(fn(sigma))
      case Proj2(sigma) => aggregator_(fn(sigma))
      case IdType(x, y) => aggregator_(fn(x), fn(y))
      case Refl(x) => aggregator_(fn(x))
      case J(typeFunc, reflFunc) => aggregator_(fn(typeFunc), fn(reflFunc))
      case BoolType => BoolType | fn | (a => aggregator_(a))
      case BoolTrue => BoolTrue | fn | (a => aggregator_(a))
      case BoolFalse => BoolFalse | fn | (a => aggregator_(a))
      case BoolElim(ifTrue, ifFalse) => aggregator_(fn(ifTrue), fn(ifFalse))
      case BoolRec(a, b, c) => aggregator_(a | fn, b | fn, c | fn)
      case EmptyType => EmptyType | fn | (a => aggregator_(a))
      case EmptyElim(anyType) => aggregator_(fn(anyType))
      case WType(arityFn) => aggregator_(arityFn | fn)
      case WElim(e, fn_) => aggregator_(e | fn, fn_ | fn)
      case WSup(wType, label, fn_) => aggregator_(wType | fn, label | fn, fn_ | fn)
      case Var(name) => ???
      case NamedType(name, typeT) => ???
    }
  }

  def hasBruijn(expr: Ast, forbidden: Int): Boolean = {
    def q(x: Ast*): Boolean = x.exists(a => hasBruijn(a, forbidden))
    expr match {
      case UnitValue => false
      case UnitType => false
      case Universe(i) => false
      case Bruijn(i, n) => i == forbidden
      case PiType(typeT, value) => hasBruijn(typeT, forbidden) || hasBruijn(value, forbidden + 1)
      case PiValue(typeT, value) => hasBruijn(typeT, forbidden) || hasBruijn(value, forbidden + 1)
      case PiElim(piValue, argument) => hasBruijn(piValue, forbidden) || hasBruijn(argument, forbidden)
      case SumType(aType, bType) => hasBruijn(aType, forbidden)
      case SumInl(aValue, bType) => q(aValue, bType)
      case SumInr(aType, bValue) => q(aType, bValue)
      case SumElim(aFn, bFn) => q(aFn, bFn)
      case SigmaType(pi) => q(pi)
      case SigmaValue(x, aX, a) => q(x, aX, a)
      case SigmaElim(pi) => q(pi)
      case Proj1(sigma) => q(sigma)
      case Proj2(sigma) => q(sigma)
      case IdType(x, y) => q(x, y)
      case Refl(x) => q(x)
      case J(typeFunc, reflFunc) => q(typeFunc, reflFunc)
      case BoolType => false
      case BoolTrue => false
      case BoolFalse => false
      case BoolElim(ifTrue, ifFalse) => q(ifFalse, ifTrue)
      case BoolRec(a, b, c) => q(a, b, c)
      case EmptyType => false
      case EmptyElim(anyType) => q(anyType)
      case WType(arityFn) => q(arityFn)
      case WElim(e, fn_) => q(e, fn_)
      case WSup(wType, label, fn_) => q(wType, label, fn_)
      case Var(name) => ???
      case NamedType(name, typeT) => ???
    }
  }

  def prettyPrint(expr: Ast): String = expr match {
    case any => any.toString
  }
}
